#include <stdio.h>
#include <stdlib.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <unistd.h>
#include <signal.h>
#include <sys/wait.h>
#include <time.h>
#include <sys/un.h>
#include <errno.h>
#include <arpa/inet.h>

#define NUM_OF_CLIENTS 10
#define DEFAULT_PATH_STREAM "/tmp/us"
#define DEFAULT_PATH_DGRAM "/tmp/ud"
#define TROJAN_STRING "trojan"
#define DGRAM_BUFSIZE 50000
#define SELF_IP_ADDRESS "127.0.0.1"

#define SOCK_TCP 1
#define SOCK_UDP 2
#define SOCK_UDSTREAM 3
#define SOCK_UDSDGRAM 4


/**
 * This function handles the errors - currently prints the error message
 * and exits with (-1).
 * */
void handle_error(char* s) {
    printf("Error (%d): %s\n", errno, s);
    exit(-1);
}

/**
 * This function is called only when the parent gets SIGCHLD signal.
 * SIGCHLD is sent to parent when a child dies.
 * The function calls waitpid for any child process, with NO-HANG flag (exits
 * immediately, and returns when waitpid found a dead child.
 * */
void sigchld_handler(int s) {
    int pid=-20;
    /* -1 means wait for any child (not a specific child process id), WNOHANG
     * means return immediately - do not block. */
    do {
        if (pid != -20) {
            printf("Child with pid %d had just died.. RIP\n",pid);
            pid = -20;
        }
        pid = waitpid(-1, NULL, WNOHANG);
    }   while(pid > 0);

}

/**
 * This function is a combination of select() and recv().
 * It gets the same arguments as recv() with addition to one more - int secs,
 * which indicates the timeout we should set to the select() function. It calls
 * handle_error if timeout was reached and therefore should only be called when
 * we're sure a packet is coming (by protocol), and returns the recv()
 * return code if a timeout was not reached in select().
 * */
int my_recv(int sock, char *buf, int buflen, int flags, int secs) {
    struct timeval timeout;
    fd_set readfds;
    int status;

    /* select() until it returns with something other than a EINTR error code
     * (which is one we should ignore) */
    do {
        timeout.tv_sec = secs;
        timeout.tv_usec = 0;
        FD_ZERO(&readfds);
        FD_SET(sock, &readfds);
        status = select(sock+1, &readfds, NULL, NULL, &timeout);
    } while ((status < 0)&&(errno == EINTR));

    if ((status == 0) || (!FD_ISSET(sock, &readfds)))
        handle_error("select() timed-out in my_recv(), something is wrong.\n");

    return recv(sock, buf, buflen, flags);
}

/**
 * This function is the same as my_recv() only using recvfrom().
 * */
int my_recvfrom(int sock, char *buf, int buflen, int flags, struct sockaddr * their_addr, int *sin_size, int secs) {
    struct timeval timeout;
    fd_set readfds;
    int status;

    do {
        timeout.tv_sec = secs;
        timeout.tv_usec = 0;
        FD_ZERO(&readfds);
        FD_SET(sock, &readfds);
        status = select(sock+1, &readfds, NULL, NULL, &timeout);
    } while ((status < 0)&&(errno == EINTR));

    if (status == 0)
        handle_error("select() timed-out in my_recvfrom(), something is wrong.\n");
    if (!FD_ISSET(sock, &readfds))
        handle_error("my_recvfrom() error: sock is not ready for reading.\n");

    return recvfrom(sock, buf, buflen, flags, their_addr, sin_size);
}

/**
 * Adds the string that starts in address b, to the buffer that starts in string
 * in position *string_len
 * */
char *addStringToBuffer(char *string, const char *b, int *string_len) {
    strcpy(&string[*string_len], b);
    (*string_len) += strlen(b);
    if (*string_len > DGRAM_BUFSIZE)
        handle_error("buffer exceeded DATAGRAM BUFSIZE! Exiting..\n");
    return string;
}

/**
 * Adds the int b to the string in place *retbuflen, and forwards *retbuflen
 * by 4 (sizeof(int)).
 * */
char *addIntToBuffer(char *string, int b, int *retbuflen) {
    int i;
    for (i=0; i<4; i++)
        string[(*retbuflen)+i] = *(((char*)&b)+i);
    *retbuflen += 4;
    if (*retbuflen > DGRAM_BUFSIZE)
        handle_error("buffer exceeded DATAGRAM BUFSIZE! Exiting..\n");
    return string;
}

/**
 * Retrieves an integer from buffer[*bufposition].
 **/
int getIntFromBuffer(char *buffer, int *bufposition) {
    int ret = 0;
    int i = 0;
    /* I know this is ugly, but the pretty bit-shifting method did not work. */
    for (i=0; i<4; i++)
        *(((char*)&ret)+i) = buffer[(*bufposition)+i];
    *bufposition += 4;
    return ret;
}

/**
 * Retrieves a string of length retbuflen from buffer[*bufposition] and
 * returns it in retbuf.
 **/
int getStringFromBuffer(char *buffer, int *bufposition, char *retbuf, int retbuflen) {
    if ((*bufposition+retbuflen) > DGRAM_BUFSIZE)
        handle_error("read excess memory from DATAGRAM BUFFER.\n");
    memcpy(retbuf, &buffer[*bufposition], retbuflen*sizeof(char));
    *bufposition += retbuflen;
    return 0;
}

/**
 * This is the main routine for the fork()ed processes. It reads the data from
 * the stream socket, builds a datagram and sends it to the server.
 **/
int childprocess(int sock, int socktype, int port) {
    struct sockaddr_in inetaddr;
    struct sockaddr_un unixaddr;
    struct sockaddr *genericsockaddr=NULL;
    int genericsock;
    int addrlen=0;
    int buflen=0;
    char *buf=NULL;
    char *retbuf=NULL;
    int retbuflen=0;
    int cnt=0,i;
    
    retbuf = malloc(sizeof(char)*DGRAM_BUFSIZE);
    if (0 == retbuf) {
        printf("error malloc()ing the retbuf.\n");
        return (EXIT_FAILURE);
    }
    memset(retbuf, '\0', DGRAM_BUFSIZE);
    retbuflen = 0;

    /* 4 first bytes are 0xffffffff (== -1 to represent
     * there's an address coming */
    retbuf = addIntToBuffer(retbuf, -1, &retbuflen);
    
    if (socktype == SOCK_TCP) {
        genericsock = socket(AF_INET, SOCK_DGRAM, 0);
        if (genericsock < 0) {
            printf("error initializing socket() for child process.\n");
            return (EXIT_FAILURE);
        }
        addrlen = sizeof(struct sockaddr_in);
        if (-1 == getpeername(sock, (struct sockaddr *)&inetaddr, &addrlen)) {
            printf("getpeername() error.\n");
            return (EXIT_FAILURE);
        }
        /* Add the address to the buffer */
        retbuf = addIntToBuffer(retbuf, inetaddr.sin_addr.s_addr, &retbuflen);
        /* And prepare inetaddr for use later in the sendto() */
        inetaddr.sin_addr.s_addr = inet_addr(SELF_IP_ADDRESS);
        inetaddr.sin_port = htons(port);
        inetaddr.sin_family = AF_INET;
        addrlen = sizeof(inetaddr);        
        genericsockaddr = (struct sockaddr *)&inetaddr;
    }
    else if (socktype == SOCK_UDSTREAM) {
        genericsock = socket(AF_UNIX, SOCK_DGRAM, 0);
        if (genericsock < 0) {
            printf("error initializing socket() for child process.\n");
            return (EXIT_FAILURE);
        }
        addrlen = sizeof(struct sockaddr_un);
        if (-1 == getpeername(sock, (struct sockaddr *)&unixaddr, &addrlen)) {
            printf("getpeername() error.\n");
            return (EXIT_FAILURE);
        }
        buflen = strlen(unixaddr.sun_path);
        /* Adding the string length */
        retbuf = addIntToBuffer(retbuf, buflen, &retbuflen);
        /* And the string (address) itself. */
        retbuf = addStringToBuffer(retbuf, unixaddr.sun_path, &retbuflen);
        unixaddr.sun_family = AF_UNIX;
        strcpy(unixaddr.sun_path, DEFAULT_PATH_DGRAM);
        addrlen = strlen(unixaddr.sun_path)+sizeof(unixaddr.sun_family);
        genericsockaddr = (struct sockaddr *)&unixaddr;
    }

    if (-1 == my_recv(sock, (char *)&cnt, 4, 0, 20)) {
        printf("error recv()ing number of processes.\n");
        return (EXIT_FAILURE);
    }

    retbuf = addIntToBuffer(retbuf, cnt, &retbuflen);

    for (i=0; i<cnt; i++) {
        /* Getting process length: */
        if (-1 == my_recv(sock, (char *)&buflen, 4, 0, 20)) {
            printf("error recv()ing length of process.\n");
            return (EXIT_FAILURE);
        }
        retbuf = addIntToBuffer(retbuf, buflen, &retbuflen);
        /* Allocating sufficient space for it */
        if (0 == (buf = malloc(sizeof(char)*(buflen+1)))) {
            printf("malloc() error when trying to get process name.\n");
            return (EXIT_FAILURE);
        }
        /* Receiving it: */
        if (buflen > 0)
            if (-1 == my_recv(sock, buf, buflen, 0, 20)) {
                printf("error (%d) recv()ing process name.\n", errno);
                return (EXIT_FAILURE);
            }
        buf[buflen] = '\0';
        /* And writing it to the buffer: */
        retbuf = addStringToBuffer(retbuf, buf, &retbuflen);
        free(buf);
    }


    /* We finished assembling our dgram and it is ready to be sent: */
    if (-1 == (i = sendto(genericsock, retbuf, DGRAM_BUFSIZE, 0, genericsockaddr, addrlen))) {
        printf("error (%d) sendto()ing data to father (mother?) process.\n", errno);
        return (EXIT_FAILURE);
    }

    shutdown(genericsock, 2);
    shutdown(sock, 2);
    /* Father exits */
    return (EXIT_SUCCESS);
}

int getMaxOfFour(int a, int b, int c, int d) {
    int ret = a;
    if (b>ret)
        ret = b;
    if (c>ret)
        ret = c;
    if (d>ret)
        ret = d;
    return ret;
}


int main(int argc, char** argv) {
    struct sockaddr_in inetaddr, their_inet_addr;
    struct sockaddr_un unixaddrstream, unixaddrdgram, their_unix_addr;
    int tcplisten, udplisten, udsstream, udsdgram;
    fd_set readfds;
    struct timeval timeout;
    int status;
    int newsock;
    unsigned int sin_size = sizeof(struct sockaddr_in);
    unsigned int sun_size_dgram, sun_size = sizeof(struct sockaddr_un);
    struct sigaction sa;
    char *buf, *tmpaddrchar, *dgrambuffer;
    int buflen=0, dgrambufferpos=0;
    int cnt = 0, i=0, found=0;
    unsigned int tmpaddrint=0;
    int port = 0;
    
    if (argc != 2)
        handle_error("Syntax error: server [port]\n");
    
    port = atoi(argv[1]);
    
    inetaddr.sin_family = AF_INET;
    inetaddr.sin_port = htons(port);
    inetaddr.sin_addr.s_addr = htonl(INADDR_ANY);

    unixaddrstream.sun_family = AF_UNIX;
    strcpy(unixaddrstream.sun_path, DEFAULT_PATH_STREAM);

    unixaddrdgram.sun_family = AF_UNIX;
    strcpy(unixaddrdgram.sun_path, DEFAULT_PATH_DGRAM);

    unlink(unixaddrstream.sun_path);
    unlink(unixaddrdgram.sun_path);
    sun_size = strlen(unixaddrstream.sun_path) + sizeof(unixaddrstream.sun_family);
    sun_size_dgram = strlen(unixaddrdgram.sun_path) + sizeof(unixaddrdgram.sun_family);

    /* Telling the OS that SIGCHLD signals should go to sigchld_handler() */
    sa.sa_handler = sigchld_handler; /* reap all dead processes */
    sigemptyset(&sa.sa_mask);
    sa.sa_flags = SA_RESTART;
    if (sigaction(SIGCHLD, &sa, NULL) == -1)
        handle_error("sigaction() error.");

    tcplisten = socket(AF_INET, SOCK_STREAM, 0);
    udplisten = socket(AF_INET, SOCK_DGRAM, 0);
    udsstream = socket(AF_UNIX, SOCK_STREAM, 0);
    udsdgram = socket(AF_UNIX, SOCK_DGRAM, 0);

    if (tcplisten<0)
        handle_error("socket() error on initializing TCP sock.\n");

    if (udplisten<0)
        handle_error("socket() error initializing UDP sock.\n");

    if (udsstream<0)
        handle_error("socket() error initializing UDS STREAM sock.\n");

    if (udsdgram<0)
        handle_error("socket() error initializing UDS DGRAM sock.\n");

    if (-1 == bind(tcplisten, (struct sockaddr *)&inetaddr, sizeof(inetaddr)))
        handle_error("bind() error on TCP sock.\n");

    if (-1 == bind(udplisten, (struct sockaddr *)&inetaddr, sizeof(inetaddr)))
        handle_error("bind() error on UDP sock.\n");

    if (-1 == bind(udsdgram, (struct sockaddr *)&unixaddrdgram, sizeof(unixaddrdgram)))
        handle_error("bind() error on UDS DGRAM sock.\n");

    if (-1 == bind(udsstream, (struct sockaddr *)&unixaddrstream, sizeof(unixaddrstream)))
        handle_error("bind() error on UDS STREAM sock.\n");

    if (-1 == listen(tcplisten, NUM_OF_CLIENTS))
        handle_error("listen() error on TCP.\n");
    
    if (-1 == listen(udsstream, NUM_OF_CLIENTS))
        handle_error("listen() error on UDS.\n");

    /* The main loop for the server- select() and decide what to do: */
    while (1) {
        /* These are suggested to be re-set after each invokation: */
        FD_ZERO(&readfds);
        FD_SET(tcplisten, &readfds);
        FD_SET(udplisten, &readfds);
        FD_SET(udsstream, &readfds);
        FD_SET(udsdgram, &readfds);
        timeout.tv_sec = 20;
        timeout.tv_usec = 0;

        if (-1 == (status = select(getMaxOfFour(tcplisten,udplisten,udsstream,udsdgram)+1,
                                    &readfds, NULL, NULL, &timeout))) {
            if (errno != EINTR) /* EINTR means some signal (like the ones we get
                 * from the children) has been received before any FD was ready.
                 * So all pages on the internet say ignore this error. */
                handle_error("select() error.");
            else
                continue;
        }
        else if (status > 0) { /* then some socket is ready */
            if (FD_ISSET(tcplisten, &readfds)) {
                printf("INET STREAM is ready.\n");
                newsock = accept(tcplisten, (struct sockaddr *)&their_inet_addr, &sin_size);
                if (newsock < 0)
                    handle_error("accept() error on TCP.");
                if (!fork()) {
                    exit(childprocess(newsock, SOCK_TCP, port));
                }
            }

            if (FD_ISSET(udsstream, &readfds)) {
                printf("UDS STREAM is ready.\n");
                newsock = accept(udsstream, (struct sockaddr *)&their_unix_addr, &sun_size);
                if (newsock < 0)
                    handle_error("accept() error on UDS STREAM.");
                if (!fork()) {
                    exit(childprocess(newsock, SOCK_UDSTREAM, port));
                }
            }

            if (FD_ISSET(udplisten, &readfds)) {
                printf("INET DGRAM is ready\n");
                tmpaddrint = 0;
                found = 0;

                dgrambuffer = (char*)malloc(DGRAM_BUFSIZE * sizeof(char));
                if (dgrambuffer == 0)
                    handle_error("malloc() error when allocating buffer for dgram.\n");
                memset(dgrambuffer, '\0', DGRAM_BUFSIZE*sizeof(char));
                
                /* get the datagram: */
                if (-1 == (status = my_recvfrom(udplisten, dgrambuffer, DGRAM_BUFSIZE, 0, (struct sockaddr *)&their_inet_addr, &sin_size, 20)))
                    handle_error("recvfrom() error in DGRAM.\n");

                dgrambufferpos = 0;

                /* get number of processes on that machine: */
                cnt = getIntFromBuffer(dgrambuffer, &dgrambufferpos);


                /* if that number is -1 this means that this is a local dgram */
                if (cnt == -1) {
                    /* get the real address: */
                    tmpaddrint = (unsigned int)getIntFromBuffer(dgrambuffer, &dgrambufferpos);
                    /* and get the real count of processes: */
                    cnt = getIntFromBuffer(dgrambuffer, &dgrambufferpos);
                }
                for (i=0;i<cnt;i++) { /* for each process: */
                    /* get process name length */
                    buflen = getIntFromBuffer(dgrambuffer, &dgrambufferpos);
                    /* allocate room for that process name */
                    if (0 == (buf = malloc(sizeof(char)*(buflen+1))))
                        handle_error("malloc() error.\n");
                    /* recv the process name */
                    getStringFromBuffer(dgrambuffer, &dgrambufferpos, buf, buflen);
                    buf[buflen] = '\0'; /* add the null-termination */
                    if (NULL != strstr(buf, TROJAN_STRING))
                        found = 1;
                    free(buf);
                }

                if (tmpaddrint != 0)
                        their_inet_addr.sin_addr.s_addr = tmpaddrint;

                if (found == 1)
                    printf("Found trojan process on: %s.\n",inet_ntoa(their_inet_addr.sin_addr));
                else
                    printf("Did not find trojan process on: %s.\n",inet_ntoa(their_inet_addr.sin_addr));

                free(dgrambuffer);
            }

            if (FD_ISSET(udsdgram, &readfds)) {
                tmpaddrint = 0;
                tmpaddrchar = NULL;
                found = 0;
                printf("UDS DGRAM is ready.\n");

                dgrambuffer = (char*)malloc(DGRAM_BUFSIZE * sizeof(char));
                memset(dgrambuffer, '\0', DGRAM_BUFSIZE*sizeof(char));
                if (dgrambuffer == 0)
                    handle_error("malloc() error when allocating buffer for dgram.\n");

                /* get the datagram: */
                if (-1 == (status = my_recvfrom(udsdgram, dgrambuffer, DGRAM_BUFSIZE, 0, (struct sockaddr *)&their_unix_addr, &sun_size_dgram, 20)))
                    handle_error("recvfrom() error in DGRAM.\n");

                dgrambufferpos = 0;

                /* get number of processes on that machine: */
                cnt = getIntFromBuffer(dgrambuffer, &dgrambufferpos);

                /* if that number is -1 this means that this is a local dgram */
                if (cnt == -1) {
                    /* get the real address: */
                    tmpaddrint = getIntFromBuffer(dgrambuffer, &dgrambufferpos);
                    if (0 == (tmpaddrchar = malloc(sizeof(char)*(tmpaddrint+1))))
                        handle_error("malloc() error.\n");
                    getStringFromBuffer(dgrambuffer, &dgrambufferpos, tmpaddrchar, tmpaddrint);
                    tmpaddrchar[tmpaddrint] = '\0';

                    /* and get the real count of processes: */
                    cnt = getIntFromBuffer(dgrambuffer, &dgrambufferpos);
                }
                for (i=0;i<cnt;i++) { /* for each process: */
                    /* get process name length */
                    buflen = getIntFromBuffer(dgrambuffer, &dgrambufferpos);

                    /* allocate room for that process name */
                    if (0 == (buf = malloc(sizeof(char)*(buflen+1))))
                        handle_error("malloc() error.\n");

                    /* recv the process name */
                    getStringFromBuffer(dgrambuffer, &dgrambufferpos, buf, buflen);
                    buf[buflen] = '\0'; /* add the null-termination */
                    if (NULL != strstr(buf, TROJAN_STRING))
                        found = 1;
                    free(buf);
                }

                if (tmpaddrchar != NULL) {
                    if (found == 1)
                        printf("Found trojan process on: UDS STREAM.\n");
                    else
                        printf("Did not find trojan process on: UDS STREAM.\n");
                    free(tmpaddrchar);
                } else {
                    if (found == 1)
                        printf("Found trojan process on: UDS DGRAM.\n");
                    else
                        printf("Did not find trojan process on: UDS DGRAM.\n");
                }
                free(dgrambuffer);
            }

            continue;
        }
        /* got here if select timed-out */
    }

    return (EXIT_SUCCESS);
}